import cv2
import numpy as np
import os
import csv
import shutil  # 导入 shutil 模块用于文件夹操作

class CellFeatureExtraction:
    def __init__(self, folder_path, output_file):
        self.folder_path = folder_path
        self.output_file = output_file

    def shape_irregularity(self, contour):
        """
        计算红细胞的形状不规则度
        :param contour: 输入的轮廓，为 numpy 数组，形状为 (n, 2) 或 (n, 1, 2)
        :return: 形状不规则度
        """
        if contour.ndim == 3:  # 检查轮廓维度是否为 (n, 1, 2)
            contour = contour[:, 0, :]
        contour_complex = np.empty(contour.shape[:1], dtype=complex)
        contour_complex.real = contour[:, 0]
        contour_complex.imag = contour[:, 1]
        fourier_transform = np.fft.fft(contour_complex)
        num_coefficients = 10
        low_freq_coefficients = np.copy(fourier_transform)
        low_freq_coefficients[num_coefficients:-num_coefficients] = 0
        inverse_transform = np.fft.ifft(low_freq_coefficients)
        mse = np.mean((contour_complex - inverse_transform) ** 2)
        return mse

    def extract_cell_features(self, image_path):
        # 读取图像
        image = cv2.imread(image_path)
        if image is None:
            print ("无法读取图像")
            return None

        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # 查找轮廓
        contours, _ = cv2.findContours(gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        if not contours:
            print(f"未在图像中找到轮廓: {image_path}")
            return None

        # 只处理最大的轮廓
        max_contour = max(contours, key=cv2.contourArea)

        # 计算面积和周长
        area = cv2.contourArea(max_contour)
        perimeter = cv2.arcLength(max_contour, True)

        # 计算质点
        M = cv2.moments(max_contour)
        if M['m00'] != 0:
            cx = int(M['m10'] / M['m00'])
            cy = int(M['m01'] / M['m00'])
        else:
            cx, cy = 0, 0

        # 使用新方法计算长轴和短轴
        if len(max_contour) >= 5:
            ellipse = cv2.fitEllipse(max_contour)
            (center, axes, orientation) = ellipse
            major_axis_length = max(axes)
            minor_axis_length = min(axes)

            # 计算其他特征
            max_diameter = area / (4*perimeter)
            circularity_factor = major_axis_length / minor_axis_length if minor_axis_length != 0 else 0
            form_factor = (4 * np.pi * area) / (perimeter ** 2) if perimeter != 0 else 0
            deviation_factor = (circularity_factor / area) * 10000 if area != 0 else 0
            semilunarity_factor = major_axis_length / (major_axis_length - minor_axis_length) if (major_axis_length - minor_axis_length) != 0 else 0

            # 计算形状不规则度
            irregularity = abs(self.shape_irregularity(max_contour))

            return {
                'area': area,
                'perimeter': perimeter,
                'max_diameter': max_diameter,
                'major_axis': major_axis_length,
                'minor_axis': minor_axis_length,
                'circularity_factor': circularity_factor,
                'form_factor': form_factor,
                'deviation_factor': deviation_factor,
                'semilunarity_factor': semilunarity_factor,
                'centroid_x': cx,
                'centroid_y': cy,
                'irregularity': irregularity
            }
        else:
            print(f"轮廓点不足以拟合椭圆: {image_path}")
            return None

    def process_folder(self):
        results = []
        for filename in os.listdir(self.folder_path):
            if filename.endswith(('.png', '.jpg', '.jpeg', '.bmp')):
                image_path = os.path.join(self.folder_path, filename)
                try:
                    features = self.extract_cell_features(image_path)
                    if features:
                        features['filename'] = filename
                        results.append(features)
                except Exception as e:
                    print(f"处理图像时出错 {filename}: {str(e)}")
        return results

    def save_to_csv(self, results):
        os.makedirs(os.path.dirname(self.output_file), exist_ok=True)
        with open(self.output_file, 'w', newline='') as csvfile:
            fieldnames = ['filename', 'area', 'perimeter', 'max_diameter', 'major_axis',
                          'minor_axis', 'circularity_factor', 'form_factor',
                          'semilunarity_factor', 'deviation_factor', 'irregularity',
                          'centroid_x', 'centroid_y']
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

            writer.writeheader()
            for cell in results:
                writer.writerow(cell)

        print(f"结果已保存到 {self.output_file}")
        return self.output_file

    def run(self):
        results = self.process_folder()
        if results:
            csv_path = self.save_to_csv(results)
            # 移除处理过的文件夹
            try:
                shutil.rmtree(self.folder_path)
                print(f"已移除文件夹: {self.folder_path}")
            except Exception as e:
                print(f"移除文件夹时出错: {str(e)}")
            
            return csv_path
        else:
            print("未提取到有效特征")
            return None